# search_server.py
import os
from typing import List
from mcp.server.fastmcp import FastMCP
from openai import OpenAI
from sentence_transformers import SentenceTransformer
import chromadb
import io
import sys

# 在文件开头添加编码设置
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

# Create server
mcp = FastMCP("Search Server")

# 千问模型配置
QWEN_MODEL = "qwen-turbo"
QWEN_API_KEY = "sk-4e88cf4db3e14894bafaff606d296610"
QWEN_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"

# 本地向量数据库
chroma_client = None
collection = None
embedding_model = None

def initialize_chroma_db():
    """初始化本地Chroma向量数据库"""
    global chroma_client, collection, embedding_model

    print("加载嵌入模型...")
    try:
        embedding_model = SentenceTransformer(r'D:\ideaSpace\MyPython\models\all-MiniLM-L6-v2')
        print("嵌入模型加载成功")
    except Exception as e:
        print(f"嵌入模型加载失败: {e}")
        embedding_model = None

    try:
        chroma_client = chromadb.PersistentClient(path="./chroma_db")
        print("Chroma客户端初始化成功")

        collection = chroma_client.get_or_create_collection(
            name="documents",
            metadata={"hnsw:space": "cosine"}
        )
        print("集合创建成功")

    except Exception as e:
        print(f"Chroma数据库初始化失败: {e}")
        raise

def call_qwen_api(messages: list, model: str = QWEN_MODEL) -> str:
    """调用千问API"""
    try:
        client = OpenAI(
            api_key=QWEN_API_KEY,
            base_url=QWEN_BASE_URL
        )

        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0.7
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"千问API调用失败: {e}")
        return f"API调用错误: {str(e)}"

def _run_rag(query: str) -> str:
    """使用本地向量数据库进行RAG搜索"""
    if collection is None:
        return "向量数据库未初始化"

    try:
        if embedding_model is not None:
            query_embedding = embedding_model.encode([query])[0].tolist()
            results = collection.query(
                query_embeddings=[query_embedding],
                n_results=3
            )
        else:
            results = collection.query(
                query_texts=[query],
                n_results=3
            )

        if results['documents'] and results['documents'][0]:
            return "\n\n".join(results['documents'][0])
        else:
            return "未找到相关文档"

    except Exception as e:
        print(f"RAG搜索错误: {e}")
        return f"搜索错误: {str(e)}"

def _summarize_rag_response(rag_output: str) -> str:
    """使用千问模型进行摘要"""
    if "错误" in rag_output or "未找到" in rag_output:
        return rag_output

    messages = [
        {"role": "system", "content": "你是一个专业的文本摘要助手，请用中文简洁地总结文本。"},
        {"role": "user", "content": f"请用中文简洁地总结以下文本：\n\n{rag_output}"}
    ]

    return call_qwen_api(messages)

@mcp.tool()
def generate_rag_output(query: str) -> str:
    """生成RAG输出"""
    print(f"[RAG] 查询: {query}")
    rag_output = _run_rag(query)
    print(f"[RAG] 结果长度: {len(rag_output)}")
    return _summarize_rag_response(rag_output)

@mcp.tool()
def run_web_search(query: str) -> str:
    """运行网络搜索"""
    print(f"[搜索] 查询: {query}")

    messages = [
        {"role": "system", "content": "你是一个专业的AI助手，能够回答各种问题。请用中文回答。"},
        {"role": "user", "content": query}
    ]

    return call_qwen_api(messages)

@mcp.tool()
def list_indexed_documents() -> List[str]:
    """列出已索引的文档"""
    if collection is None:
        return ["向量数据库未初始化"]

    try:
        count = collection.count()
        return [f"已索引 {count} 个文档"]
    except Exception as e:
        return [f"获取文档列表失败: {str(e)}"]

def index_documents(directory: str):
    """索引文档到本地向量数据库"""
    if not os.path.exists(directory):
        print(f"目录不存在: {directory}")
        return False

    SUPPORTED_EXTENSIONS = {'.txt', '.md'}
    supported_files = []

    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)
        if os.path.isfile(file_path):
            _, ext = os.path.splitext(filename)
            if ext.lower() in SUPPORTED_EXTENSIONS:
                supported_files.append(file_path)

    if not supported_files:
        print("没有找到支持的文件")
        return False

    print(f"开始索引 {len(supported_files)} 个文件...")

    documents = []
    embeddings = []
    ids = []

    for i, file_path in enumerate(supported_files):
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read().strip()

            if not content:
                continue

            if embedding_model is not None:
                embedding = embedding_model.encode([content])[0].tolist()
                embeddings.append(embedding)

            documents.append(content)
            ids.append(f"doc_{i}")

            print(f"已索引: {os.path.basename(file_path)}")

        except Exception as e:
            print(f"索引文件失败 {file_path}: {e}")

    if documents:
        try:
            if embeddings and len(embeddings) == len(documents):
                collection.add(
                    documents=documents,
                    embeddings=embeddings,
                    ids=ids
                )
            else:
                collection.add(
                    documents=documents,
                    ids=ids
                )

            print(f"成功索引 {len(documents)} 个文档")
            return True

        except Exception as e:
            print(f"添加到向量数据库失败: {e}")
            return False
    return False

if __name__ == "__main__":
    # 检查千问API密钥
    qwen_api_key = os.environ.get("DASHSCOPE_API_KEY")
    if not qwen_api_key:
        print("警告: 未设置 DASHSCOPE_API_KEY 环境变量")

    print("初始化本地向量数据库...")
    try:
        initialize_chroma_db()
    except Exception as e:
        print(f"向量数据库初始化失败: {e}")
        print("将继续运行，但RAG功能可能受限")

    current_dir = os.path.dirname(os.path.abspath(__file__))
    samples_dir = os.path.join(current_dir, "sample_files")

    if not os.path.exists(samples_dir):
        os.makedirs(samples_dir)
        print(f"创建了样本目录: {samples_dir}")

    # 索引文档
    success = index_documents(samples_dir)

    print("启动MCP服务器...")

    try:
        import uvicorn
        from fastapi import FastAPI, HTTPException, Request
        from fastapi.middleware.cors import CORSMiddleware
        from fastapi.responses import StreamingResponse
        from mcp.server.sse import SseServerTransport
        import anyio
        import json

        print("✅ uvicorn 和 fastapi 导入成功")

        # 创建FastAPI应用
        app = FastAPI(title="MCP Search Server", version="1.0.0")

        # 添加CORS中间件
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

        # 创建SSE传输实例
        sse_transport = SseServerTransport("/sse")

        @app.get("/sse")
        async def sse_endpoint2(request: Request):
            sse_transport.connect_sse(request.scope, request.receive, request._send)

        # 添加SSE端点 - 新版本方式
        @app.post("/sse")
        async def sse_endpoint(request: Request):
            """处理SSE连接"""
            try:
                # 读取请求体
                body = await request.body()
                if body:
                    try:
                        message = json.loads(body.decode('utf-8'))
                        print(f"📨 收到SSE消息: {message}")
                    except json.JSONDecodeError:
                        print("⚠️  无法解析SSE消息体")

                # 创建SSE响应
                async def event_stream():
                    try:
                        # 这里需要实现SSE流处理
                        # 在实际应用中，这里应该处理MCP协议消息
                        yield "data: {}\n\n".format(json.dumps({"type": "connected"}))

                        # 模拟一些事件
                        await anyio.sleep(1)
                        yield "data: {}\n\n".format(json.dumps({"type": "ready"}))

                    except Exception as e:
                        print(f"SSE流错误: {e}")

                return StreamingResponse(
                    event_stream(),
                    media_type="text/event-stream",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                    }
                )

            except Exception as e:
                print(f"SSE端点错误: {e}")
                raise HTTPException(500, "SSE连接失败")

        # 添加工具调用端点
        @app.post("/tools/call")
        async def call_tool_endpoint(payload: dict):
            """调用MCP工具"""
            try:
                tool_name = payload.get("name")
                arguments = payload.get("arguments", {})

                print(f"📞 调用工具: {tool_name}, 参数: {arguments}")

                # 根据工具名调用相应的函数
                if tool_name == "generate_rag_output":
                    query = arguments.get("query", "")
                    if not query:
                        raise HTTPException(400, "缺少query参数")
                    result = generate_rag_output(query)

                elif tool_name == "run_web_search":
                    query = arguments.get("query", "")
                    if not query:
                        raise HTTPException(400, "缺少query参数")
                    result = run_web_search(query)

                elif tool_name == "list_indexed_documents":
                    result = list_indexed_documents()

                else:
                    raise HTTPException(404, f"工具 '{tool_name}' 不存在")

                return {"success": True, "result": result, "tool": tool_name}

            except HTTPException:
                raise
            except Exception as e:
                print(f"工具调用错误: {e}")
                raise HTTPException(500, f"工具调用失败: {str(e)}")

        # 添加工具列表端点
        @app.get("/tools/list")
        async def list_tools_endpoint():
            """获取可用工具列表"""
            tools = [
                {
                    "name": "generate_rag_output",
                    "description": "生成RAG输出",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "用户查询"
                            }
                        },
                        "required": ["query"]
                    }
                },
                {
                    "name": "run_web_search",
                    "description": "运行网络搜索",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "用户查询"
                            }
                        },
                        "required": ["query"]
                    }
                },
                {
                    "name": "list_indexed_documents",
                    "description": "列出已索引的文档",
                    "inputSchema": {
                        "type": "object",
                        "properties": {}
                    }
                }
            ]
            return {"tools": tools}

        # 添加健康检查端点
        @app.get("/health")
        async def health_check():
            """健康检查"""
            return {
                "status": "healthy",
                "service": "mcp-search-server",
                "version": "1.0.0"
            }

        # 添加根端点
        @app.get("/")
        async def root():
            """根端点"""
            return {
                "message": "MCP Search Server",
                "endpoints": {
                    "sse": "POST /sse",
                    "tools_list": "GET /tools/list",
                    "tools_call": "POST /tools/call",
                    "health": "GET /health"
                }
            }

        print("🚀 HTTP服务器启动成功!")
        print("📍 访问地址: http://localhost:8000")
        print("🔗 SSE端点: POST http://localhost:8000/sse")
        print("🛠️  工具列表: GET http://localhost:8000/tools/list")
        print("📞 工具调用: POST http://localhost:8000/tools/call")

        uvicorn.run(app, host="localhost", port=8000, log_level="info")

    except ImportError as e:
        print(f"❌ 依赖导入失败: {e}")
        print("💡 使用纯SSE模式")
        print("🔗 SSE端点: http://localhost:8000/sse")

        # 纯SSE模式
        mcp.run(transport="sse")





"""@mcp.tool()解释：
    装饰器工厂：调用 mcp.tool() 返回一个装饰器函数
    类方法装饰器：用于装饰类的方法（虽然MCP工具通常是模块级函数）
    元数据收集器：收集函数的名称、参数、文档等信息
    注册器：将工具注册到MCP服务器中
tool函数签名分解：
    def tool(
        self,                    # ① 实例方法，属于某个类
        name: str | None = None, # ② 工具名称参数
        title: str | None = None, # ③ 工具标题参数  
        description: str | None = None, # ④ 工具描述参数
        annotations: ToolAnnotations | None = None, # ⑤ 工具注解参数
        structured_output: bool | None = None, # ⑥ 结构化输出参数
    ) -> Callable[[AnyFunction], AnyFunction]: # ⑦ 返回值类型
等价的手动调用：
    # 装饰器语法
    @mcp.tool(name="my_tool")
    def my_function():
        pass
    
    # 等价于：
    decorator = mcp.tool(name="my_tool")  # 调用tool方法，返回装饰器函数
    my_function = decorator(my_function)  # 用装饰器包装函数
"""